#!/usr/bin/env python3
import argparse
import os
import sys
from time import sleep

import grpc

# Import P4Runtime lib from parent utils dir
# Probably there's a better way of doing this.
sys.path.append(
    os.path.join(os.path.dirname(os.path.abspath(__file__)),
                 '../../utils/'))
import p4runtime_lib.bmv2
import p4runtime_lib.helper
from p4runtime_lib.error_utils import printGrpcError
from p4runtime_lib.switch import ShutdownAllSwitchConnections

# def writeForwardRules(p4info_helper, ingress_sw, egress_sw,
#                       match_fields, action_params):
#     table_entry = p4info_helper.buildTableEntry(
#         table_name="MyIngress.ipv4_lpm",
#         match_fields={
#             "hdr.ipv4.dstAddr": match_fields
#         },
#         action_name="MyIngress.ipv4_forward",
#         action_params=action_params
#     )
#     ingress_sw.WriteTableEntry(table_entry)

def writeForwardRules(p4info_helper, ingress_sw, egress_sw, match_fields, action_params):
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ipv4_lpm",
        match_fields={"hdr.ipv4.dstAddr": match_fields},
        action_name="MyIngress.ipv4_forward",
        action_params=action_params
    )
    try:
        ingress_sw.WriteTableEntry(table_entry)
    except grpc.RpcError as e:
        # If the entry already exists, modify it instead of adding a new one
        if "ALREADY_EXISTS" in str(e):
            ingress_sw.ModifyTableEntry(table_entry)
        else:
            raise e


def writeSwitchId(p4info_helper, ingress_sw, egress_sw, action_params):
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyEgress.swid",
        default_action=True,
        action_name="MyEgress.set_swid",
        action_params=action_params
    )
    ingress_sw.WriteTableEntry(table_entry)

def main(p4info_file_path, bmv2_file_path):
    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)

    try:
        # Define switch connections
        switches = []
        for i in range(1, 5):
            switch = p4runtime_lib.bmv2.Bmv2SwitchConnection(
                name=f's{i}',
                address=f'127.0.0.1:5005{i}',
                device_id=i - 1,
                proto_dump_file=f'logs/s{i}-p4runtime-requests.txt'
            )
            switch.MasterArbitrationUpdate()
            switch.SetForwardingPipelineConfig(
                p4info=p4info_helper.p4info,
                bmv2_json_file_path=bmv2_file_path
            )
            print(f"Installed P4 Program using SetForwardingPipelineConfig on s{i}")
            switches.append(switch)

        # Define switch IDs
        switch_ids = [1, 2, 3, 4]

        # Define forwarding rules and switch IDs
        forwarding_rules = [
            {"dstAddr": "08:00:00:00:01:11", "port": 1},  # s1
            {"dstAddr": "08:00:00:00:02:22", "port": 2},  # s1
            {"dstAddr": "08:00:00:00:03:00", "port": 3},  # s1
            {"dstAddr": "08:00:00:00:04:00", "port": 4},  # s1
            {"dstAddr": "08:00:00:00:03:00", "port": 4},  # s2
            {"dstAddr": "08:00:00:00:04:00", "port": 3},  # s2
            {"dstAddr": "08:00:00:00:03:33", "port": 1},  # s2
            {"dstAddr": "08:00:00:00:04:44", "port": 2},  # s2
            {"dstAddr": "08:00:00:00:01:00", "port": 1},  # s3
            {"dstAddr": "08:00:00:00:01:00", "port": 1},  # s3
            {"dstAddr": "08:00:00:00:02:00", "port": 2},  # s3
            {"dstAddr": "08:00:00:00:02:00", "port": 2},  # s3
            {"dstAddr": "08:00:00:00:01:00", "port": 2},  # s4
            {"dstAddr": "08:00:00:00:01:00", "port": 2},  # s4
            {"dstAddr": "08:00:00:00:02:00", "port": 1},  # s4
            {"dstAddr": "08:00:00:00:02:00", "port": 1}   # s4
        ]

        # Write forwarding rules and switch IDs
        for i, rule in enumerate(forwarding_rules):
            # print(switches[i // 4],rule)
            writeForwardRules(
                p4info_helper, switches[i // 4], switches[i // 4], [f"10.0.{i%4 + 1}.{i%4 + 1}", 32], rule
            )

        for i, switch_id in enumerate(switch_ids):
            writeSwitchId(
                p4info_helper, switches[i], switches[i], {"swid": switch_id}
            )

        # Keep the script running
        # while True:
        #     sleep(1)

    except KeyboardInterrupt:
        print("Shutting down.")
    except grpc.RpcError as e:
        printGrpcError(e)

    ShutdownAllSwitchConnections()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='P4Runtime Controller')
    parser.add_argument('--p4info', help='p4info proto in text format from p4c',
                        type=str, action="store", required=False,
                        default='./build/link_monitor.p4.p4info.txt')
    parser.add_argument('--bmv2-json', help='BMv2 JSON file from p4c',
                        type=str, action="store", required=False,
                        default='./build/link_monitor.json')
    args = parser.parse_args()

    if not os.path.exists(args.p4info):
        parser.print_help()
        print("\np4info file not found: %s\nHave you run 'make'?" % args.p4info)
        parser.exit(1)
    if not os.path.exists(args.bmv2_json):
        parser.print_help()
        print("\nBMv2 JSON file not found: %s\nHave you run 'make'?" % args.bmv2_json)
        parser.exit(1)
    main(args.p4info, args.bmv2_json)
